package com.sudo.rpc.server.task;

import com.sudo.rpc.api.utils.RpcRes;
import com.sudo.rpc.server.annotation.Note;
import com.sudo.rpc.server.task.registry.BaseTaskRegistry;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.Map;

/**
 * @Author：DayDayUp
 * @Create：2023/2/15 16:40
 */
@Slf4j
public class ShellModule extends BaseTaskRegistry {

    private final static Runtime runtime = Runtime.getRuntime();

    public ShellModule(){
        setModuleClassMap(getClass());
    }

    private final static String[] invalidWords = {" su ", " sudo "};
    private final static String testShell = "ls .";

    /**
     * 执行shell
     * @param paramMap 服务节点
     */
    @Note(desc = "执行shell",params = "shell(String,true,echo 123,shell命令)")
    public RpcRes execShell(Map<String,Object> paramMap) throws IOException, InterruptedException {
        String shell = paramMap.get("shell").toString();
        if (!StringUtils.equals(shell, testShell)) {
            return RpcRes.errorMsg("测试环境仅允许执行：[ " + testShell +" ]命令");
        }
        if (StringUtils.containsAny(shell,invalidWords)) {
            return RpcRes.errorMsg("shell中不可包含su、sudo等敏感命令");
        }
        String[] shellArr = new String[]{"/bin/sh", "-c",shell};
        String result = getProcessResult(shellArr);
        log.info("执行结果：{}", result);
        return RpcRes.okMsg("执行结果：" + result);
    }

    private String getProcessResult(String[] shellArr) throws IOException, InterruptedException {
        Process process = runtime.exec(shellArr);
        int exitValue = process.waitFor();
        if (0 != exitValue) {
            log.info("call shell failed. error code is: " + exitValue);
            throw new RuntimeException("call shell failed. error code is: " + exitValue);
        }
        BufferedReader br = new BufferedReader(new InputStreamReader(process.getInputStream()));
        StringBuilder stringBuffer = new StringBuilder();
        String line;
        while ((line = br.readLine()) != null) {
            stringBuffer.append(line).append("\n");
        }
        br.close();
        return stringBuffer.toString();
    }
}
